#include <Windows.h>
#include "KinectLib.h"
#include "Utils.h"

KinectLib::KinectLib()
{
    mCallback           = NULL;
    mNuiSensor          = NULL;
    
    mColorEnabled       = false;
    mColorEvent         = NULL;
    mColorHandle        = NULL;
    mColorThread        = NULL;
    mColorImageSize.cx  = 0;
    mColorImageSize.cy  = 0;

    mDepthEnabled       = false;
    mDepthEvent         = NULL;
    mDepthHandle        = NULL;
    mDepthThread        = NULL;
    mDepthImageSize.cx  = 0;
    mDepthImageSize.cy  = 0;

    mSkeletonEnabled    = false;
    mSkeletonEvent      = NULL;
    mSkeletonThread     = NULL;

    mIsStop             = true;

    mRecordSkeletonFile = "";
    mIsRecordSkeleton   = false;
    mReplaySkeletonFile = "";
    mFileReplaySkeleton = NULL;
    mIsReplaySkeleton   = false;
    mIsReplayPaused     = false;
    mReplayDelayMS      = 30;
    mReplaySkeletonThread = NULL;
}

KinectLib::~KinectLib()
{
}

void KinectLib::setCallback(IKinectCallback *pCallback)
{
    mCallback = pCallback;
}

DWORD WINAPI KinectLib::_ThreadProc(LPVOID lpParam)
{
    ThreadContext *pContext = (ThreadContext *)lpParam;
    pContext->pKinect->onThreadProc(pContext->id);
    delete pContext;

    return 0;
}

DWORD KinectLib::onThreadProc(int id)
{
    if(id == 1)
    {
        NUI_IMAGE_FRAME imageFrame = { 0 };

        while(!mIsStop)
        {
            DWORD dwRet = WaitForSingleObject(mColorEvent, INFINITE);
            if(dwRet == WAIT_OBJECT_0)
            {
                if(mIsStop)
                {
                    break ;
                }
                HRESULT hr = mNuiSensor->NuiImageStreamGetNextFrame(mColorHandle, 0, &imageFrame);
                if(FAILED(hr))
                {
                    continue ;
                }

                INuiFrameTexture * pTexture = imageFrame.pFrameTexture;
                NUI_LOCKED_RECT lockedRect = { 0 };
                pTexture->LockRect(0, &lockedRect, NULL, 0);                
                if(lockedRect.Pitch != 0)
                {
                    if(mCallback != NULL)
                    {
                        KinectImageFrame frame = { 0 };
                        frame.width     = mColorImageSize.cx;
                        frame.height    = mColorImageSize.cy;
                        frame.stride    = lockedRect.Pitch;
                        frame.bytes     = lockedRect.size;
                        frame.bits      = lockedRect.pBits;
                        mCallback->onKinectColorFrame(this, frame);
                    }
                }
                pTexture->UnlockRect(0);
                mNuiSensor->NuiImageStreamReleaseFrame(mColorHandle, &imageFrame);
            }
        }
    }
    else if(id == 2)
    {
        NUI_IMAGE_FRAME imageFrame = { 0 };

        while(!mIsStop)
        {
            DWORD dwRet = WaitForSingleObject(mDepthEvent, INFINITE);
            if(dwRet == WAIT_OBJECT_0)
            {
                if(mIsStop)
                {
                    break ;
                }
                HRESULT hr = mNuiSensor->NuiImageStreamGetNextFrame(mDepthHandle, 0, &imageFrame);
                if(FAILED(hr))
                {
                    continue ;
                }

                BOOL nearMode = FALSE;
                INuiFrameTexture * pTexture = NULL;
                hr = mNuiSensor->NuiImageFrameGetDepthImagePixelFrameTexture(
                    mDepthHandle, &imageFrame, &nearMode, &pTexture);
                if(SUCCEEDED(hr))
                {
                    NUI_LOCKED_RECT lockedRect = { 0 };
                    pTexture->LockRect(0, &lockedRect, NULL, 0);                
                    if(lockedRect.Pitch != 0)
                    {
                        const int minDepth = (nearMode ? NUI_IMAGE_DEPTH_MINIMUM_NEAR_MODE : 
                            NUI_IMAGE_DEPTH_MINIMUM) >> NUI_IMAGE_PLAYER_INDEX_SHIFT;
                        const int maxDepth = (nearMode ? NUI_IMAGE_DEPTH_MAXIMUM_NEAR_MODE : 
                            NUI_IMAGE_DEPTH_MAXIMUM) >> NUI_IMAGE_PLAYER_INDEX_SHIFT;
                        if(mCallback != NULL)
                        {
                            KinectDepthFrame frame = { 0 };
                            frame.width     = mDepthImageSize.cx;
                            frame.height    = mDepthImageSize.cy;
                            frame.minDepth  = minDepth;
                            frame.maxDepth  = maxDepth;
                            frame.depths    = reinterpret_cast<NUI_DEPTH_IMAGE_PIXEL *>(lockedRect.pBits);
                            mCallback->onKinectDepthFrame(this, frame);
                        }
                    }
                    pTexture->UnlockRect(0);
                    pTexture->Release();
                }
                mNuiSensor->NuiImageStreamReleaseFrame(mDepthHandle, &imageFrame);
            }
        }
    }
    else if(id == 3)
    {
        FILE * pRecordFile = NULL;
        NUI_SKELETON_FRAME skeletonFrame = { 0 };
        while(!mIsStop)
        {
            DWORD dwRet = WaitForSingleObject(mSkeletonEvent, INFINITE);
            if(dwRet == WAIT_OBJECT_0)
            {
                if(mIsStop)
                {
                    break ;
                }

                if(mIsRecordSkeleton && (pRecordFile == NULL))
                {
                    fopen_s(&pRecordFile, mRecordSkeletonFile.c_str(), "wb");
                }

                if((!mIsRecordSkeleton) && (pRecordFile != NULL))
                {
                    fclose(pRecordFile);
                    pRecordFile = NULL;
                }

                memset(&skeletonFrame, 0, sizeof(skeletonFrame));
                HRESULT hr = mNuiSensor->NuiSkeletonGetNextFrame(0, &skeletonFrame);
                hr = mNuiSensor->NuiTransformSmooth(&skeletonFrame, NULL);

                if(SUCCEEDED(hr))
                {
                    if(pRecordFile != NULL)
                    {
                        fwrite(&skeletonFrame, 1, sizeof(skeletonFrame), pRecordFile);
                    }

                    if(mCallback != NULL)
                    {
                        mCallback->onKinectSkeletonFrame(this, skeletonFrame);
                    }
                }
            }
        }
        if(pRecordFile != NULL)
        {
            fclose(pRecordFile);
            pRecordFile = NULL;
        }
    }
    else if(id == 1001)
    {
        NUI_SKELETON_FRAME skeletonFrame = { 0 };
        while(mIsReplaySkeleton)
        {
            if(!mIsReplayPaused)
            {
                int readSize = fread(&skeletonFrame, 1, sizeof(skeletonFrame), mFileReplaySkeleton);
                if(readSize < sizeof(skeletonFrame))
                {
                    fseek(mFileReplaySkeleton, 0, SEEK_SET);
                }
                if(readSize == sizeof(skeletonFrame))
                {
                    if(!KinectLib::checkSkeletonTracked(skeletonFrame))
                    {
                        continue ;
                    }

                    if(mCallback != NULL)
                    {
                        mCallback->onKinectSkeletonFrame(this, skeletonFrame);
                    }
                }
            }
            Utils::longSleep(mReplayDelayMS, mIsReplaySkeleton);
        }
    }
    return id;
}

HRESULT KinectLib::open(bool enableColor, bool enableDepth, bool enableSkeleton)
{
    INuiSensor * pNuiSensor = NULL;
    int sensorCount = 0;
    HRESULT hr = NuiGetSensorCount(&sensorCount);
    if(FAILED(hr))
    {
        return hr;
    }

    for(int i = 0; i < sensorCount; i ++)
    {
        hr = NuiCreateSensorByIndex(i, &pNuiSensor);
        if(FAILED(hr))
        {
            continue;
        }

        hr = pNuiSensor->NuiStatus();
        if (S_OK == hr)
        {
            mNuiSensor = pNuiSensor;
            break;
        }

        pNuiSensor->Release();
    }

    if(NULL == mNuiSensor)
    {
        return E_FAIL;
    }

    DWORD dwFlags = 0;
    if(enableColor)
    {
        dwFlags |= NUI_INITIALIZE_FLAG_USES_COLOR;
    }
    if(enableDepth)
    {
        dwFlags |= NUI_INITIALIZE_FLAG_USES_DEPTH_AND_PLAYER_INDEX;
    }
    if(enableSkeleton)
    {
        dwFlags |= NUI_INITIALIZE_FLAG_USES_SKELETON;
    }    
    hr = mNuiSensor->NuiInitialize(dwFlags); 
    if(FAILED(hr))
    {
        return hr;
    }

    mIsStop = false;

    if(enableColor)
    {
        mColorEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
        hr = mNuiSensor->NuiImageStreamOpen(NUI_IMAGE_TYPE_COLOR, NUI_IMAGE_RESOLUTION_640x480,
            0, 2, mColorEvent, &mColorHandle);
        mColorImageSize.cx = 640;
        mColorImageSize.cy = 480;
        mColorThread = startThread(1);
    }

    if(enableDepth)
    {
        mDepthEvent = CreateEvent(NULL, TRUE, FALSE, NULL);
        hr = mNuiSensor->NuiImageStreamOpen(NUI_IMAGE_TYPE_DEPTH_AND_PLAYER_INDEX, NUI_IMAGE_RESOLUTION_640x480,
            0, 2, mDepthEvent, &mDepthHandle);
        mDepthImageSize.cx = 640;
        mDepthImageSize.cy = 480;
        mDepthThread = startThread(2);
    }

//    NuiCreateInteractionStream();

    if(enableSkeleton)
    {
        mSkeletonEvent = CreateEventW(NULL, TRUE, FALSE, NULL);
        hr = mNuiSensor->NuiSkeletonTrackingEnable(mSkeletonEvent, 0);

        mSkeletonThread = startThread(3);
    }

    return hr;
}

HANDLE KinectLib::startThread(int id)
{
    ThreadContext *pContext = new ThreadContext();
    pContext->pKinect = this;
    pContext->id = id;
    DWORD threadId = 0;
    HANDLE hThread = CreateThread(NULL, 0, KinectLib::_ThreadProc, pContext, 0, &threadId);
    if(hThread == NULL)
    {
        delete pContext;
    }
    return hThread;
}

void KinectLib::stopThreads()
{
    mIsStop = true;
    if(mColorThread != NULL)
    {
        SetEvent(mColorEvent);

        WaitForSingleObject(mColorThread, INFINITE);
        CloseHandle(mColorThread);
        mColorThread = NULL;
    }
    if(mDepthThread != NULL)
    {
        SetEvent(mDepthEvent);

        WaitForSingleObject(mDepthThread, INFINITE);
        CloseHandle(mDepthThread);
        mDepthThread = NULL;
    }
    if(mSkeletonThread != NULL)
    {
        SetEvent(mSkeletonEvent);

        WaitForSingleObject(mSkeletonThread, INFINITE);
        CloseHandle(mSkeletonThread);
        mSkeletonThread = NULL;
    }
}

void KinectLib::close()
{
    stopRecordSkeleton();
    stopReplaySkeleton();

    stopThreads();

    if(mNuiSensor != NULL)
    {
        mNuiSensor->NuiShutdown();
        mNuiSensor->Release();
        mNuiSensor = NULL;
    }
    if((mColorHandle != NULL) && (mColorHandle != INVALID_HANDLE_VALUE))
    {
        mColorHandle = NULL;
    }
    if((mDepthHandle != NULL) && (mDepthHandle != INVALID_HANDLE_VALUE))
    {
        mDepthHandle = NULL;
    }
    if((mSkeletonEvent != NULL) && (mSkeletonEvent != INVALID_HANDLE_VALUE))
    {
        CloseHandle(mSkeletonEvent);
        mSkeletonEvent = NULL;
    }
}

SkeletonPoint KinectLib::convertSkeletonPosition(const Vector4 &position, int width, int height)
{
    LONG x = 0;
    LONG y = 0;
    USHORT depth = 0;

    NuiTransformSkeletonToDepthImage(position, &x, &y, &depth, NUI_IMAGE_RESOLUTION_640x480);

    SkeletonPoint result;
    result.x        = (x*width)/640;
    result.y        = (y*height)/480;
    result.depth    = (int)(position.z*1000.0f);
    return result;
}

bool KinectLib::checkSkeletonTracked(const NUI_SKELETON_FRAME &frame)
{
    for(int s = 0 ; s < NUI_SKELETON_COUNT; s ++)
    {
        const NUI_SKELETON_TRACKING_STATE trackingState = frame.SkeletonData[s].eTrackingState;
        if(trackingState == NUI_SKELETON_TRACKED)
        {
            return true;
        }
    }
    return false;
}

// for debugging
void KinectLib::setRecordSkeletonFile(const char *skeleonFile)
{
    mRecordSkeletonFile = (skeleonFile != NULL) ? skeleonFile : "";
}

void KinectLib::startRecordSkeleton()
{
    mIsRecordSkeleton = true;
}

void KinectLib::stopRecordSkeleton()
{
    mIsRecordSkeleton = false;
}

bool KinectLib::isRecordSkeletonStarted()
{
    return mIsRecordSkeleton;
}

void KinectLib::setReplaySkeletonFile(const char *skeleonFile)
{
    mReplaySkeletonFile = (skeleonFile != NULL) ? skeleonFile : "";
}

void KinectLib::setReplaySkeletonDelay(int delayMS)
{
    mReplayDelayMS = delayMS;
}

int KinectLib::getReplaySkeletonDelay()
{
    return mReplayDelayMS;
}

bool KinectLib::startReplaySkeleton()
{
    if(mFileReplaySkeleton != NULL)
    {
        fclose(mFileReplaySkeleton);
        mFileReplaySkeleton = NULL;
    }

    fopen_s(&mFileReplaySkeleton, mReplaySkeletonFile.c_str(), "rb");
    if(mFileReplaySkeleton == NULL)
    {
        return false;
    }

    mIsReplaySkeleton = true;
    mReplaySkeletonThread = startThread(1001);
    return (mReplaySkeletonThread != NULL);
}

void KinectLib::stopReplaySkeleton()
{
    mIsReplaySkeleton = false;
    if(mReplaySkeletonThread != NULL)
    {
        WaitForSingleObject(mReplaySkeletonThread, INFINITE);
        CloseHandle(mReplaySkeletonThread);
        mReplaySkeletonThread = NULL;
    }

    if(mFileReplaySkeleton != NULL)
    {
        fclose(mFileReplaySkeleton);
        mFileReplaySkeleton = NULL;
    }
}

bool KinectLib::isReplaySkeletoStarted()
{
    return mIsReplaySkeleton;
}

void KinectLib::pauseOrResumeReplaySkeleton()
{
    mIsReplayPaused = !mIsReplayPaused;
}

bool KinectLib::isReplaySkeletonPaused()
{
    return mIsReplayPaused;
}